import cv2
import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk
import numpy as np
from enum import Enum

class TrackingAlgorithm(Enum):
    MEANSHIFT = 1
    CAMSHIFT = 2

class VideoPlayer:
    def __init__(self, root):
        """
        初始化视频播放器
        
        Args:
            root: Tkinter根窗口
        """
        self.root = root
        self.root.title("视频播放器")
        # 设置窗口最小尺寸
        self.root.minsize(800, 600)
        
        # 视频相关变量
        self.video_path = ""
        self.cap = None
        self.is_playing = False
        self.current_frame = 0
        self.captured_image = None  # 存储截取的图像
        
        # 矩形选择相关变量
        self.selecting_roi = False
        self.roi_start = None
        self.roi_end = None
        self.roi_rect = None
        
        # 跟踪相关变量
        self.tracking_algorithm = TrackingAlgorithm.MEANSHIFT
        self.tracking_enabled = False
        self.track_window = None
        self.roi_hist = None
        # 使用与示例相同的终止条件
        self.term_crit = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 1)
        
        # 用于存储帧尺寸信息和ROI信息
        self.frame_width = 0
        self.frame_height = 0
        self.selected_roi = None  # 存储选择的ROI区域信息 (x, y, w, h)
        
        # 创建界面
        self.create_widgets()
    
    def create_widgets(self):
        """创建GUI组件"""
        # 主框架
        self.main_frame = tk.Frame(self.root)
        self.main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 上方框架 - 视频显示和截图显示
        self.top_frame = tk.Frame(self.main_frame)
        self.top_frame.pack(fill=tk.BOTH, expand=True)
        
        # 视频显示区域(使用Canvas支持绘制)
        self.video_canvas = tk.Canvas(self.top_frame, width=640, height=480, bg="black")
        self.video_canvas.pack(side=tk.LEFT, padx=5, pady=5, fill=tk.BOTH, expand=True)
        self.video_image = None  # 用于存储Canvas上的图像引用
        
        # 绑定鼠标事件
        self.video_canvas.bind("<Button-1>", self.on_mouse_down)
        self.video_canvas.bind("<B1-Motion>", self.on_mouse_drag)
        self.video_canvas.bind("<ButtonRelease-1>", self.on_mouse_up)
        
        # 截取图像显示区域
        self.imageview = tk.Label(self.top_frame, text="截取图像将显示在这里", 
                                 width=20, height=15, bg="lightgray", relief=tk.RAISED)
        self.imageview.pack(side=tk.RIGHT, padx=5, pady=5)
        
        # 控制按钮框架
        self.control_frame = tk.Frame(self.main_frame)
        self.control_frame.pack(fill=tk.X, pady=5)
        
        # 基本控制按钮
        self.open_btn = tk.Button(self.control_frame, text="打开文件", command=self.open_file, width=8)
        self.open_btn.pack(side=tk.LEFT, padx=2)
        
        self.play_btn = tk.Button(self.control_frame, text="播放", command=self.play_video, width=8)
        self.play_btn.pack(side=tk.LEFT, padx=2)
        
        self.pause_btn = tk.Button(self.control_frame, text="暂停", command=self.pause_video, width=8)
        self.pause_btn.pack(side=tk.LEFT, padx=2)
        
        self.stop_btn = tk.Button(self.control_frame, text="停止", command=self.stop_video, width=8)
        self.stop_btn.pack(side=tk.LEFT, padx=2)
        
        # 跟踪控制按钮
        self.track_btn = tk.Button(self.control_frame, text="开始跟踪", command=self.enable_tracking, width=8)
        self.track_btn.pack(side=tk.LEFT, padx=2)
        
        # 算法选择
        self.algorithm_label = tk.Label(self.control_frame, text="算法:")
        self.algorithm_label.pack(side=tk.LEFT, padx=(10, 2))
        
        self.algorithm_var = tk.StringVar(value="meanshift")
        self.algorithm_menu = tk.OptionMenu(self.control_frame, self.algorithm_var, "meanshift", "camshift")
        self.algorithm_menu.config(width=10)
        self.algorithm_menu.pack(side=tk.LEFT, padx=2)
        
        # 状态栏
        self.status_frame = tk.Frame(self.main_frame)
        self.status_frame.pack(fill=tk.X, pady=5)
        
        self.status_label = tk.Label(self.status_frame, text="就绪", relief=tk.SUNKEN, anchor=tk.W)
        self.status_label.pack(fill=tk.X)
    
    def on_mouse_down(self, event):
        """鼠标按下事件"""
        if not self.is_playing and self.cap and self.cap.isOpened():
            self.selecting_roi = True
            self.roi_start = (event.x, event.y)
            self.roi_end = (event.x, event.y)
            
            # 删除之前的矩形（如果存在）
            if self.roi_rect:
                self.video_canvas.delete(self.roi_rect)
    
    def on_mouse_drag(self, event):
        """鼠标拖动事件"""
        if self.selecting_roi:
            self.roi_end = (event.x, event.y)
            # 删除之前的矩形
            if self.roi_rect:
                self.video_canvas.delete(self.roi_rect)
            # 绘制新的矩形
            self.roi_rect = self.video_canvas.create_rectangle(
                self.roi_start[0], self.roi_start[1],
                self.roi_end[0], self.roi_end[1],
                outline="red", width=2
            )
    
    def on_mouse_up(self, event):
        """鼠标释放事件"""
        if self.selecting_roi and self.roi_start and self.roi_end:
            self.selecting_roi = False
            end_x, end_y = event.x, event.y
            
            # 确保矩形有效（至少5x5像素）
            if abs(end_x - self.roi_start[0]) > 5 and abs(end_y - self.roi_start[1]) > 5:
                # 删除临时矩形
                if self.roi_rect:
                    self.video_canvas.delete(self.roi_rect)
                    self.roi_rect = None
                
                # 获取当前帧
                ret, frame = self.cap.read()
                if ret:
                    # 重置到当前帧
                    self.cap.set(cv2.CAP_PROP_POS_FRAMES, self.current_frame)
                    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    
                    # 计算实际图像中的坐标
                    canvas_width = max(self.video_canvas.winfo_width(), 1)
                    canvas_height = max(self.video_canvas.winfo_height(), 1)
                    
                    if frame is not None and hasattr(frame, 'shape'):
                        self.frame_height, self.frame_width = frame.shape[:2]
                        x_ratio = self.frame_width / canvas_width
                        y_ratio = self.frame_height / canvas_height
                    else:
                        self.update_status("无法获取帧信息")
                        return
                    
                    # 转换坐标
                    x1 = int(min(self.roi_start[0], end_x) * x_ratio)
                    y1 = int(min(self.roi_start[1], end_y) * y_ratio)
                    x2 = int(max(self.roi_start[0], end_x) * x_ratio)
                    y2 = int(max(self.roi_start[1], end_y) * y_ratio)
                    
                    # 确保坐标在图像范围内
                    x1 = max(0, x1)
                    y1 = max(0, y1)
                    x2 = min(self.frame_width, x2)
                    y2 = min(self.frame_height, y2)
                    
                    # 确保区域有效
                    if x2 > x1 and y2 > y1:
                        # 截取区域
                        cropped = frame[y1:y2, x1:x2]
                        img = Image.fromarray(cropped)
                        self.captured_image = ImageTk.PhotoImage(image=img)
                        self.imageview.config(image=self.captured_image)
                        
                        # 保存选择的ROI区域信息，但不立即开始跟踪
                        self.selected_roi = (x1, y1, x2-x1, y2-y1)
                        self.update_status(f"已选择ROI区域: {self.selected_roi}，点击'开始跟踪'按钮启动跟踪")
            
            # 清理矩形
            if self.roi_rect:
                self.video_canvas.delete(self.roi_rect)
                self.roi_rect = None
    
    def open_file(self):
        """打开视频文件并显示第一帧"""
        self.video_path = filedialog.askopenfilename(filetypes=[("视频文件", "*.mp4 *.avi *.mov")])
        if self.video_path:
            self.stop_video()
            self.cap = cv2.VideoCapture(self.video_path)
            
            if not self.cap.isOpened():
                self.update_status("无法打开视频文件")
                return
                
            self.current_frame = 0
            
            # 提取并显示第一帧
            ret, frame = self.cap.read()
            if ret:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                # 保存帧尺寸信息
                self.frame_height, self.frame_width = frame.shape[:2]
                
                img = Image.fromarray(frame)
                imgtk = ImageTk.PhotoImage(image=img)
                
                # 如果已有图像，先删除
                if self.video_image:
                    self.video_canvas.delete(self.video_image)
                
                # 创建新图像
                self.video_image = self.video_canvas.create_image(
                    self.video_canvas.winfo_width() // 2 or self.frame_width // 2, 
                    self.video_canvas.winfo_height() // 2 or self.frame_height // 2, 
                    anchor=tk.CENTER, 
                    image=imgtk
                )
                self.video_canvas.image = imgtk  # 保持引用
                
                # 重置视频到第一帧
                self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                self.update_status("视频已加载")
    
    def play_video(self):
        """播放视频"""
        if self.cap and not self.is_playing:
            self.is_playing = True
            self.update_status("正在播放")
            self.update_frame()
    
    def pause_video(self):
        """暂停视频"""
        self.is_playing = False
        self.update_status("已暂停")
    
    def stop_video(self):
        """停止视频"""
        self.is_playing = False
        if self.cap:
            self.cap.release()
            self.cap = None
        self.current_frame = 0
        if self.video_image:
            self.video_canvas.delete(self.video_image)
            self.video_image = None
        self.imageview.config(image="", text="截取图像将显示在这里")
        # 重置跟踪状态
        self.tracking_enabled = False
        self.track_window = None
        self.roi_hist = None
        self.frame_width = 0
        self.frame_height = 0
        self.selected_roi = None
        self.update_status("已停止")
    
    def enable_tracking(self):
        """启用跟踪功能"""
        # 只有在选择了ROI区域后才能开始跟踪
        if self.selected_roi is not None:
            self.tracking_enabled = True
            algorithm = self.algorithm_var.get()
            if algorithm == "meanshift":
                self.tracking_algorithm = TrackingAlgorithm.MEANSHIFT
            elif algorithm == "camshift":
                self.tracking_algorithm = TrackingAlgorithm.CAMSHIFT
            
            # 使用已选择的ROI区域初始化跟踪器
            x, y, w, h = self.selected_roi
            self.start_tracking(x, y, w, h)
            self.update_status(f"已开始跟踪，使用算法: {algorithm}")
        else:
            self.update_status("请先选择要跟踪的区域")
    
    def start_tracking(self, x, y, w, h):
        """初始化跟踪器"""
        if not self.cap or self.frame_width == 0 or self.frame_height == 0:
            self.update_status("无法初始化跟踪器")
            return
        
        # 获取当前帧（不改变当前帧位置）
        self.cap.set(cv2.CAP_PROP_POS_FRAMES, self.current_frame)
        ret, frame = self.cap.read()
        if not ret:
            self.update_status("无法读取当前帧")
            return
        
        # 确保ROI在图像范围内
        x = max(0, min(x, self.frame_width - 1))
        y = max(0, min(y, self.frame_height - 1))
        w = max(1, min(w, self.frame_width - x))
        h = max(1, min(h, self.frame_height - y))
        
        # 检查ROI区域是否有效
        if w <= 0 or h <= 0:
            self.update_status("无效的ROI区域")
            return
            
        # 设置ROI
        roi = frame[y:y+h, x:x+w]
        # 使用与示例相同的色彩空间转换
        hsv_roi = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)
        # 使用与示例相同的直方图计算方法
        roi_hist = cv2.calcHist([hsv_roi], [0], None, [180], [0, 180])
        cv2.normalize(roi_hist, roi_hist, 0, 255, cv2.NORM_MINMAX)
        
        # 保存直方图和跟踪窗口
        self.roi_hist = roi_hist
        self.track_window = (x, y, w, h)
        self.update_status(f"跟踪初始化完成: x={x}, y={y}, w={w}, h={h}")
    
    def update_tracking(self, frame):
        """更新跟踪结果"""
        if not self.tracking_enabled or self.track_window is None or self.roi_hist is None:
            return frame
        
        # 检查frame是否有效
        if frame is None or not hasattr(frame, 'shape'):
            return frame
            
        # 使用与示例相同的色彩空间转换和反向投影计算
        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
        dst = cv2.calcBackProject([hsv], [0], self.roi_hist, [0, 180], 1)
        
        try:
            if self.tracking_algorithm == TrackingAlgorithm.MEANSHIFT:
                # 使用与示例相同的meanShift调用方式
                ret, self.track_window = cv2.meanShift(dst, self.track_window, self.term_crit)
            elif self.tracking_algorithm == TrackingAlgorithm.CAMSHIFT:
                # 使用与示例类似的CamShift调用方式
                ret, self.track_window = cv2.CamShift(dst, self.track_window, self.term_crit)
            
            # 绘制跟踪结果
            x, y, w, h = self.track_window
            # 确保绘制坐标在图像范围内
            x = max(0, x)
            y = max(0, y)
            w = max(0, min(w, self.frame_width - x))
            h = max(0, min(h, self.frame_height - y))
            
            # 使用与示例相同的绘制方式
            cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
        except cv2.error as e:
            # 如果跟踪失败，禁用跟踪
            self.update_status(f"跟踪错误: {e}")
            self.tracking_enabled = False
            
        return frame
    
    def update_frame(self):
        """更新视频帧"""
        if self.is_playing and self.cap and self.cap.isOpened():
            ret, frame = self.cap.read()
            if ret:
                # 更新跟踪结果
                if self.tracking_enabled and self.track_window is not None:
                    frame = self.update_tracking(frame)
                
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                img = Image.fromarray(frame)
                imgtk = ImageTk.PhotoImage(image=img)
                
                # 如果已有图像，先删除
                if self.video_image:
                    self.video_canvas.delete(self.video_image)
                
                # 创建新图像
                self.video_image = self.video_canvas.create_image(
                    self.video_canvas.winfo_width() // 2 or self.frame_width // 2, 
                    self.video_canvas.winfo_height() // 2 or self.frame_height // 2, 
                    anchor=tk.CENTER, 
                    image=imgtk
                )
                self.video_canvas.image = imgtk  # 保持引用
                self.current_frame += 1
                self.root.after(30, self.update_frame)
            else:
                # 视频播放结束后自动回到开头继续播放
                self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                self.update_frame()
    
    def update_status(self, message):
        """更新状态栏信息"""
        self.status_label.config(text=message)
        print(message)  # 同时输出到控制台

if __name__ == "__main__":
    root = tk.Tk()
    app = VideoPlayer(root)
    root.mainloop()